package documentsampler;

import informationExtractionSystem.InformationExtractor;

import java.util.ArrayList;
import java.util.HashSet;

import utils.FileHandlerUtils;

import distributedRE.model.DistributedREModel;
import distributedRE.model.Document;
import distributedRE.model.Relation;
import distributedRE.model.Session;

public class PoolSampler implements DocumentSampler {

	private long sample_size;
	private HashSet<Document> all;
	private ArrayList<Document> bad;
	private ArrayList<Document> good;
	private boolean generated;
	
	public PoolSampler(long sample_size){
		this.sample_size = sample_size;
		this.generated = false;
	}
	
	public ArrayList<Document> getBadSample(DistributedREModel dREModel){

		if (generated){
			generated = false;
			good = new ArrayList<Document>();
			return bad;
		}
		
		build(dREModel);
		
		return bad;
		
	}

	private void build(DistributedREModel dREModel) {
		
		generated = true;
		
		ArrayList<Document> badones = new ArrayList<Document>();
		ArrayList<Document> goodones = new ArrayList<Document>();
		
		createAll(dREModel);
		
		InformationExtractor ie = dREModel.getInformationExtractor();
		Relation r = dREModel.getRelation();
		
		for (Document document : all) {
			
			if (ie.isUseful(r,document) == InformationExtractor.USELESS_VALUE){
				badones.add(document);
			}else if (ie.isUseful(r,document) == InformationExtractor.USEFUL_VALUE){
				goodones.add(document);
			}
			
		}
		
		Integer[] indexes = FileHandlerUtils.generateRandomIndexes(badones.size(), (int)sample_size);
		
		bad = new ArrayList<Document>();
		
		for (int i = 0; i < indexes.length; i++) {
			
			bad.add(badones.get(indexes[i].intValue()));
			
		}
		
		Integer[] indexes2 = FileHandlerUtils.generateRandomIndexes(goodones.size(), (int)sample_size);
		
		good = new ArrayList<Document>();
		
		for (int i = 0; i < indexes2.length; i++) {
			
			good.add(goodones.get(indexes2[i].intValue()));
			
		}
		
		all.clear();
		
	}

	private void createAll(DistributedREModel dREModel) {
		
		all = new HashSet<Document>();
		
		ArrayList<Session> sessions = dREModel.getAllSessions();
		
		for (Session session : sessions) {
			
			ArrayList<ArrayList<Document>> docu = session.getDocumentLists();
			
			for (ArrayList<Document> arrayList : docu) {
				
				for (Document document : arrayList) {
					
					all.add(document);
									
				}
				
			}
			
		}
		
	}

	public ArrayList<Document> getGoodSample(DistributedREModel dREModel){

		if (generated){
			generated = false;
			bad = new ArrayList<Document>();
			return good;
		}
		
		build(dREModel);
		
		return good;
	}

}
